import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
import numpy as np
import cv2 as cv
from ultralytics import YOLO
import time
from gobang_interfaces.srv import Play

# 摄像头矫正参数
k = [-0.281831359625869, 0.072276754058603, -0.007778520548340]
p = [-2.949785592530397e-04, -2.575583259892622e-04]
f = [9.796854694029854e+02, 9.797733387300718e+02]
c = [9.512391471508578e+02, 6.992317763795083e+02]

mtx = np.array([[f[0], 0.00000000e+00, c[0]],
                [0.00000000e+00, f[1], c[1]],
                [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]])
dist = np.array([k[0], k[1], p[0], p[1], k[2]])
model = YOLO("/home/haichi/dev_ws/src/gobang_vision/gobang_vision/best.pt")  # 加载YOLO模型
corner_points = None
grid_size = 19  # 棋盘格大小
last_valid_positions = []  # 保存上次有效的棋子位置（统一存储所有棋子）
all_pieces = []  # 所有棋子位置（不区分黑白）
perspective_matrix = None  # 保存投影变换矩阵

# 多次识别确认机制相关变量
CONFIRMATION_COUNT = 3  # 需要连续识别几次相同结果才确认
detection_history = []  # 存储最近几次的检测结果
confirmed_positions = []  # 已确认的有效位置（统一存储）


class ImageSubscriber(Node):
    def __init__(self):
        super().__init__('image_subscriber')
        self.subscription = self.create_subscription(
            CompressedImage,
            '/image',
            self.listener_callback,
            1)
        self.get_logger().info('---------------------------------1')
        self.gobang_play_service_client = self.create_client(
            Play, 'gobang_play')
        while not self.gobang_play_service_client.wait_for_service(timeout_sec=1.0):
            self.get_logger().info('Service not available, waiting again...')
        self.get_logger().info('---------------------------------2')

    def is_in_bounds(self, point):
        return point[0] - 2 >= 0 and point[1] - 2 >= 0

    def listener_callback(self, msg):
        global corner_points, perspective_matrix, all_pieces
        # 将JPEG数据解码为OpenCV格式
        np_arr = np.frombuffer(msg.data, np.uint8)
        img = cv.imdecode(np_arr, cv.IMREAD_COLOR)
        if img is not None:
            # self.get_logger().info(f'Received image of shape: {img.shape}')
            # 摄像头矫正
            h, w = img.shape[:2]
            newcameramtx, roi = cv.getOptimalNewCameraMatrix(
                mtx, dist, (w, h), 1, (w, h))
            img = cv.undistort(img, mtx, dist, None, mtx)
            display_size = (1600, 1200)
            dst_resized = cv.resize(img, display_size, interpolation=cv.INTER_AREA)
            cv.imshow('Undistorted Image', dst_resized)
            cv.waitKey(1)
            # 棋盘检测与投影变换
            if corner_points is None:
                points = find_chessboard(dst_resized)
                if points is not None:
                    corner_points = points
                    print(f'Chessboard corners detected: {corner_points}')
                else:
                    self.get_logger().error('No chessboard corners found')
                    return
            # 绘制检测到的棋盘角点
            temp = cv.polylines(dst_resized, [corner_points.astype(np.int32)], True, (0, 255, 0), 2)
            cv.imshow('Chessboard Detection', temp)
            cv.waitKey(1)

            warped_image = perform_perspective_transform(dst_resized, corner_points)
            cv.imshow('Warped Image', warped_image)
            cv.waitKey(1)
            current_pieces = detect_pieces(warped_image, model, None)
            
            # 检查是否有新的棋子落下
            if len(current_pieces) > len(all_pieces):
                # 找到新下的棋子（不在之前位置列表中的棋子）
                new_pieces = [p for p in current_pieces if p not in all_pieces]
                if new_pieces:
                    new_piece = new_pieces[0]  # 假设一次只下一颗棋子
                    all_pieces = current_pieces.copy()
                    self.get_logger().info(f'New piece detected at: {new_piece}')
                    
                    # 发送棋子位置信息
                    if self.is_in_bounds(new_piece):
                        request = Play.Request()
                        request.row = new_piece[1] - 2
                        request.col = new_piece[0] - 2
                        future = self.gobang_play_service_client.call_async(request)
                        self.get_logger().info(f'Sent piece position: row={request.row}, col={request.col}')
            # time.sleep(1)  # 延时，避免过快调用服务
        else:
            self.get_logger().error('Failed to decode image')

def main(args=None):
    rclpy.init(args=args)
    node = ImageSubscriber()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()

def validate_corner_angles(contour):
    """
    验证四边形的角度是否合理
    :param contour: 四边形角点
    :return: 是否为合理的四边形
    """
    points = contour.reshape(4, 2)

    # 计算四个角的角度
    angles = []
    for i in range(4):
        p1 = points[i]
        p2 = points[(i + 1) % 4]
        p3 = points[(i + 2) % 4]

        # 计算向量
        v1 = p1 - p2
        v2 = p3 - p2

        # 计算角度
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        cos_angle = np.clip(cos_angle, -1, 1)  # 防止数值误差
        angle = np.arccos(cos_angle) * 180 / np.pi
        angles.append(angle)

    # 检查角度是否合理（接近90度）
    for angle in angles:
        if angle < 45 or angle > 135:  # 角度应该在45-135度之间
            return False

    return True

def find_chessboard(image):
    """ 
    改进的棋盘识别函数，提高识别准确性
    """
    gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)

    # 使用高斯模糊减少噪声
    blurred = cv.GaussianBlur(gray, (5, 5), 0)

    # 使用自适应阈值处理
    adaptive_thresh = cv.adaptiveThreshold(blurred, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C,
                                            cv.THRESH_BINARY, 11, 2)

    # 使用多种边缘检测方法
    edges1 = cv.Canny(blurred, 30, 100)
    edges2 = cv.Canny(blurred, 50, 150)
    edges3 = cv.Canny(blurred, 80, 200)
    edges4 = cv.Canny(blurred, 110, 220)

    # 组合边缘检测结果
    edges = cv.bitwise_or(edges1, edges2)
    edges = cv.bitwise_or(edges, edges3)
    edges = cv.bitwise_or(edges, edges4)

    # 形态学操作，闭合断开的边缘
    kernel = np.ones((3, 3), np.uint8)
    edges = cv.morphologyEx(edges, cv.MORPH_CLOSE, kernel, iterations=1)
    cv.imshow('Edges', edges)
    cv.waitKey(1)
    # 寻找轮廓
    contours, _ = cv.findContours(edges, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)

    # 候选棋盘列表
    candidates = []

    for contour in contours:
        area = cv.contourArea(contour)
        perimeter = cv.arcLength(contour, True)

        # 过滤掉太小的轮廓
        if area < 5000 or perimeter < 200:
            continue

        # 多边形逼近
        epsilon = 0.01 * perimeter  # 调整精度
        approx = cv.approxPolyDP(contour, epsilon, True)

        # 如果是四边形
        if len(approx) == 4:
            # 检查是否为凸四边形
            if cv.isContourConvex(approx):
                # 计算四边形的长宽比
                rect = cv.minAreaRect(contour)
                width, height = rect[1]
                aspect_ratio = max(width, height) / min(width, height)

                # 棋盘应该接近正方形
                if aspect_ratio < 2.5:  # 长宽比不超过2.5
                    candidates.append({
                        'contour': approx,
                        'area': area,
                        'aspect_ratio': aspect_ratio,
                        'score': area / (aspect_ratio * 1000)  # 综合评分
                    })

    if not candidates:
        print("未检测到符合条件的棋盘轮廓！")
        return None

    # 按综合评分排序，选择最好的候选
    candidates.sort(key=lambda x: x['score'], reverse=True)

    # 进一步验证最佳候选
    best_candidate = candidates[0]
    best_contour = best_candidate['contour']

    # 验证四个角点的角度
    if validate_corner_angles(best_contour):
        print(f"找到棋盘！面积: {best_candidate['area']:.0f}, 长宽比: {best_candidate['aspect_ratio']:.2f}")
        return best_contour.reshape(4, 2)

    # 如果最佳候选不合格，尝试其他候选
    for candidate in candidates[1:]:
        if validate_corner_angles(candidate['contour']):
            print(f"找到棋盘！面积: {candidate['area']:.0f}, 长宽比: {candidate['aspect_ratio']:.2f}")
            return candidate['contour'].reshape(4, 2)

    print("未找到合适的棋盘！")
    return None

def order_points(points):
    """
    对棋盘的四个顶点进行排序：左上、右上、右下、左下
    :param points: 四个顶点坐标
    :return: 排序后的顶点坐标
    """
    rect = np.zeros((4, 2), dtype="float32")
    s = points.sum(axis=1)
    rect[0] = points[np.argmin(s)]
    rect[2] = points[np.argmax(s)]
    diff = np.diff(points, axis=1)
    rect[1] = points[np.argmin(diff)]
    rect[3] = points[np.argmax(diff)]
    return rect

def perform_perspective_transform(image, points=None):
    """
    对棋盘进行投影变换，如果points为None，则使用全局保存的变换矩阵
    :param image: 输入图像
    :param points: 棋盘顶点坐标，如果为None则使用全局
    :return: 投影变换后的图像
    """
    global perspective_matrix

    if points is not None:
        # 第一次识别，计算变换矩阵
        ordered_points = order_points(points)
        width = 400
        height = 400
        dst_points = np.array([
            [0, 0],
            [width - 1, 0],
            [width - 1, height - 1],
            [0, height - 1]
        ], dtype="float32")
        perspective_matrix = cv.getPerspectiveTransform(ordered_points, dst_points)

    # 使用保存的变换矩阵进行投影变换
    if perspective_matrix is not None:
        warped_image = cv.warpPerspective(image, perspective_matrix, (400, 400))
        return warped_image
    else:
        return None
    
def calculate_positions(image, pieces, matrix):
    """
    计算棋子在棋盘上的落点，并返回错误数量
    统一处理所有棋子，不区分黑白
    """
    positions = []
    error_count = 0

    for x, y in pieces:
        # 将像素坐标转换为棋盘坐标
        board_x = round(x / (400 / (grid_size - 1)))
        board_y = round(y / (400 / (grid_size - 1)))

        # 确保坐标在棋盘范围内
        board_x = max(0, min(grid_size - 1, board_x))
        board_y = max(0, min(grid_size - 1, board_y))

        # 计算棋盘落点的像素坐标
        expected_x = int(board_x * (400 / (grid_size - 1)))
        expected_y = int(board_y * (400 / (grid_size - 1)))

        # 计算偏移距离
        distance = np.sqrt((x - expected_x) ** 2 + (y - expected_y) ** 2)

        # 偏移阈值（可以根据需要调整）
        threshold = 10  # 单位：像素

        if distance > threshold:
            error_count += 1
            # 在图像中标记偏离的棋子
            cv.circle(image, (x, y), 10, (0, 255, 255), 2)
            cv.putText(image, "Error", (x - 20, y - 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
            # print(f"棋子放落位置出错，请重新放置！ 棋子像素坐标: ({x}, {y}), 偏移距离: {distance:.2f}px")
        else:
            # 检查位置唯一性，避免重复位置
            if (board_x, board_y) not in positions:
                positions.append((board_x, board_y))

    return positions, error_count

def detect_pieces(image, model, matrix):
    """
    使用 YOLO 模型检测棋子，统一识别所有棋子不区分黑白
    需要多次识别到相同结果才认为是有效的，否则返回上一次的有效结果
    :param image: 输入图像
    :param model: YOLO 模型
    :param matrix: 投影变换矩阵
    :return: 所有棋子在棋盘上的落点坐标列表
    """
    global last_valid_positions, detection_history, confirmed_positions

    results = model(image, verbose=False)  # 禁用输出信息以减少干扰
    all_detected_pieces = set()  # 存储所有检测到的棋子位置

    
    # 创建可视化图像的副本
    vis_image = image.copy()

    for result in results:
        for box in result.boxes:
            cls = int(box.cls)
            conf = float(box.conf)
            
            # 获取检测框坐标 (xyxy格式)
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            
            # 获取中心点坐标
            x_center, y_center, _, _ = box.xywh[0]
            x_center, y_center = int(x_center), int(y_center)
            
            # 统一处理所有类型的棋子，不区分黑白
            if cls in [0, 1]:  # 0是黑子，1是白子，但我们统一处理
                color = (0, 255, 0)  # 统一用绿色标记
                label = "Piece"
                all_detected_pieces.add((x_center, y_center))
            else:
                color = (0, 0, 255)  # 红色用于未知类别
                label = f"Unknown_{cls}"
            
            # print(f"检测到 {(all_detected_pieces)} ")
            # 绘制检测框
            cv.rectangle(vis_image, (x1, y1), (x2, y2), color, 2)
            
            # 绘制中心点
            cv.circle(vis_image, (x_center, y_center), 3, color, -1)
            
            # # 准备标签文本
            # label_text = f"{label}: {conf:.2f}"
            
            # # 计算文本大小
            # font = cv.FONT_HERSHEY_SIMPLEX
            # font_scale = 0.5
            # thickness = 1
            # (text_width, text_height), baseline = cv.getTextSize(label_text, font, font_scale, thickness)
            
            # # 绘制文本背景
            # cv.rectangle(vis_image, (x1, y1 - text_height - baseline - 5), 
            #             (x1 + text_width, y1), color, -1)
            
            # # 绘制文本
            # text_color = (255, 255, 255)
            # cv.putText(vis_image, label_text, (x1, y1 - baseline - 2), 
            #           font, font_scale, text_color, thickness)
    
    # 显示可视化结果
    cv.imshow('YOLO Detection Results', vis_image)
    cv.waitKey(10)

    # 计算棋子在棋盘上的落点，并统计错误数量
    piece_positions, error_count = calculate_positions(image, all_detected_pieces, matrix)

    # 如果错误棋子数量超过1个，直接返回上次确认的位置
    if (error_count >= 1 or 
        len(all_detected_pieces) > (len(confirmed_positions) + 1) or 
        len(all_detected_pieces) < len(confirmed_positions)):
        # print(f"⚠️ 检测到{error_count}个错误棋子或棋子数量异常，返回上次确认的结果")
        return confirmed_positions.copy()
    
    # 当前检测结果
    current_result = piece_positions.copy()
    
    # 将当前结果添加到历史记录
    detection_history.append(current_result)
    
    # 保持历史记录在指定长度内
    if len(detection_history) > CONFIRMATION_COUNT:
        detection_history.pop(0)
    
    # 检查是否有足够的历史记录进行确认
    if len(detection_history) < CONFIRMATION_COUNT:
        print(f"正在收集检测结果 ({len(detection_history)}/{CONFIRMATION_COUNT})，返回上次确认的结果")
        return confirmed_positions.copy()
    
    # 检查最近几次检测结果是否一致
    all_same = True
    first_result = detection_history[0]
    for result in detection_history[1:]:
        if set(result) != set(first_result):
            all_same = False
            break
    
    if all_same:
        # 连续多次检测到相同结果，确认为有效
        confirmed_positions = piece_positions.copy()
        # print(f"✅ 连续{CONFIRMATION_COUNT}次检测到相同结果，确认位置有效")
        print("棋子棋盘坐标：", piece_positions)
        return piece_positions
    else:
        # 检测结果不一致，返回上次确认的结果
        # print(f"⚠️ 检测结果不一致，返回上次确认的结果")
        return confirmed_positions.copy()